#!/usr/bin/env python
# encoding: UTF-8

from db_cleanup import get_tables
from Bio.PDB.PDBParser import PDBParser

import sys, os

db_name = sys.argv[1]
pdb_dir = sys.argv[2]

pdbs = [pdb.split("_")[1] for (pdb, _) in get_tables(db_name)]
p = PDBParser()

count = {}

for pdb in pdbs:
    model = p.get_structure(pdb, os.path.join(pdb_dir, pdb + ".pdb"))[0]
    for chain in model:
        for residue in chain:
            name = residue.get_resname()
            count[name] = count.get(name, 0) + 1

print count
            
# {'DRZ': 1, 'MA7': 3, 'SF4': 3, 'EOH': 2, '2PR': 4, 'SCC': 4, 'EPE': 1, 'SIN': 1, 'SAH': 8, '4PC': 1, 'LYS': 8976, 'PRO': 5224, 'PRN': 2, 'O2C': 1, '8OG': 4, 'SPM': 3, 'DCT': 4, 'DCP': 6, 'URA': 3, 'VAL': 7647, ' NA': 55, 'THR': 6117, '6MA': 12, '5HU': 1, ' NT': 1, 'IPA': 1, 'LAE': 4, 'PHE': 4980, '3C4': 1, 'TYR': 4544, '  C': 4, '  G': 2, 'NGA': 1, '  K': 2, 'CME': 4, 'P2U': 1, 'HDF': 4, ' CD': 1, 'OCS': 2, 'PCA': 4, 'ACY': 10, 'TSP': 6, ' YB': 2, 'TTP': 2, 'TCP': 2, 'DDY': 4, '2HP': 6, 'I5C': 2, 'FOX': 2, 'TRP': 1503, 'TRS': 6, 'DDG': 6, 'MPD': 5, 'ANG': 1, 'ASU': 1, ' AF': 1, '3MO': 1, 'TMP': 3, 'ILE': 7227, 'SUC': 49, 'UMP': 1, 'GLN': 4789, 'DOC': 1, '5CY': 1, 'NEA': 8, 'BZG': 1, 'GLY': 6721, 'DIO': 10, 'ACT': 6, 'PED': 6, '5CM': 13, 'PEG': 1, 'AKG': 4, 'GLU': 8668, '2DT': 2, 'ASP': 6163, 'SER': 6826, ' DI': 4, ' DG': 2392, ' DC': 2331, ' DA': 2389, '7GU': 1, 'FE2': 1, 'DGT': 1, ' DU': 3, ' DT': 2419, 'ASN': 5163, '3DR': 12, 'UDP': 6, 'PDI': 2, 'DUP': 3, 'UNL': 9, 'ADK': 1, 'ADI': 1, '7DA': 1, ' ZN': 113, ' CL': 23, 'BRU': 18, ' CA': 87, 'EDO': 74, 'ORP': 3, 'ALA': 8567, 'MET': 2207, 'LEU': 11190, 'MES': 1, '64T': 1, ' MG': 99, 'BAP': 1, 'AAB': 1, ' MN': 55, 'CIT': 3, '5IU': 4, 'GNE': 1, 'MSE': 86, 'FLC': 1, 'CYS': 1268, 'HIS': 2732, 'CYT': 1, 'SOS': 1, 'PPV': 2, '5PY': 3, 'FAD': 5, 'GOL': 117, '6OG': 7, 'A1P': 4, 'PO4': 6, 'NCO': 5, 'HOH': 116067, 'HXL': 1, 'PG6': 1, 'PG4': 1, 'SO4': 162, 'VO4': 4, 'HPD': 2, 'NRI': 1, 'BA2': 2, 'D3T': 1, '4MF': 1, 'ARG': 7417, '23T': 2, 'OTR': 2, 'OTS': 2}
